package AVLTree;

public class AVLTree {

    AVLNode root = null;

    static class AVLNode {
        int key;
        Object value;
        AVLNode left;
        AVLNode right;
        int height = 1;

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key, Object value, AVLNode left, AVLNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }


    //获取当前节点的高度
    private int height (AVLNode node) {
        return node == null ? 0 : node.height;
    }

    //更新当前的高度
    private void updateHeight (AVLNode node) {
        node.height = Integer.max(height(node.left),height(node.right)) + 1;
    }

    //平衡因子
    private int bf (AVLNode node) {
        return  height(node.left) - height(node.right);
    }

    //左旋
    private AVLNode leftRotate (AVLNode node) {
        AVLNode r = node.right;
        node.right = r.left;
        r.left = node;
        updateHeight(node);
        updateHeight(r);
        return r;
    }

    //右旋
    private AVLNode rightRotate (AVLNode node) {
        AVLNode l = node.left;
        node.left = l.right;
        l.right = node;
        updateHeight(node);
        updateHeight(l);
        return l;
    }

    //左右旋
    private AVLNode leftRightRotate (AVLNode node) {
        AVLNode l = node.left;
        node.left = leftRotate(l);
        return rightRotate(node);
    }

    //右左旋
    private AVLNode rightLeftRotate (AVLNode node) {
        AVLNode r = node.right;
        node.right = rightRotate(r);
        return leftRotate(node);
    }

    //检查节点是否失衡，重新平衡代码
    private AVLNode balance (AVLNode node) {
        if(node == null) {
            return  null;
        }
        if (bf(node) > 1 && bf(node.left) >= 0) {
            return rightRotate(node);
        } else if (bf(node) > 1 && bf(node.left) < 0) {
            return leftRightRotate(node);

        } else if (bf(node) < -1 && bf(node.right) <= 0) {
            return leftRotate(node);

        }else if (bf(node) < -1 && bf(node.right) > 0) {
            return rightLeftRotate(node);
        }
        return node;
    }

    //更新
    public AVLNode put (int key, Object value) {
        return doPut(root,key,value);
    }

    private AVLNode doPut(AVLNode node, int key, Object value) {
        if (node == null) {
            return new AVLNode(key,value);
        }
        if (node.key == key) {
            node.value = value;
            return node;
        }
        if (node.key > key) {
            node.left = doPut(node.left,key,value);
        }else {
            node.right = doPut(node.right,key,value);
        }
        updateHeight(node);
        return balance(node);
    }

    //删除
    public AVLNode remove (int key) {
        return doRemove(root,key);
    }
    private AVLNode doRemove (AVLNode node,int key) {
        if (node == null) {
            return null;
        }
        if (node.key > key) {
            node.left = doRemove(node.left,key);
        } else if (node.key < key) {
            node.right = doRemove(node.right,key);
        }else {

            if (node.left == null && node.right == null) {
                return null;
            } else if (node.right == null) {
                node = node.left;
            } else if (node.left == null) {
                node = node.right;
            }else {
                AVLNode p = node.right;
                while (p.left != null) {
                    p = p.left;
                }
                p.right = doRemove(node.right,p.key);
                p.left = node.left;
                node = p;
            }
        }
        updateHeight(node);
        return balance(node);
    }
}
